import os
import sys
import numpy as np
import pylab as plt
from params import *

DIR = os.path.join(WORKDIR, 'gowalla', 'final', 'cam_ready_results2')

data = {}
for file in os.listdir(DIR):
    name,ext = os.path.splitext(file)
    if not ext == '.txt':
        continue
    if not name.startswith('results'):
    #if not name.startswith('train_results'):
        continue
    print name
    t = name.split('_')
    #month, flag = t[1:3]
    month, flag = t[2:4]
    month = int(month)
    for line in open(os.path.join(DIR,file)):
        t = line.split()
        algo = t[0]
        prec,rec,fscore,auc = map(float,t[1:])
        data.setdefault(month,{}).setdefault(flag,{}).setdefault(algo,{}).setdefault('precision',[]).append(prec)
        data.setdefault(month,{}).setdefault(flag,{}).setdefault(algo,{}).setdefault('recall',[]).append(rec)
        data.setdefault(month,{}).setdefault(flag,{}).setdefault(algo,{}).setdefault('auc',[]).append(auc)


def avg(s):
    return float(sum(s))/len(s)

def avg2(s):
    return float(sum(i**2 for i in s))/len(s)

bars = {}
table = {}
for month in data:
    for flag in ['SOCIAL','PLACE','BOTH']:
        #print ''
        #print 'Month %d, dataset %s'%(month,flag)
        res = []
        for algo in data[month][flag]:
            vv = data[month][flag][algo]['precision']
            prec = sum(vv)/len(vv)
            vv = data[month][flag][algo]['recall']
            rec = sum(vv)/len(vv)
            vv = data[month][flag][algo]['auc']
            auc = avg(vv)
            auc2 = avg2(vv)
            std_auc = (auc2 - auc**2)**0.5
            print auc, std_auc
            bars.setdefault(algo,{}).setdefault(month,{})[flag]= (auc,prec,rec,std_auc)
            table.setdefault(algo,{}).setdefault(flag,{}).setdefault('auc',[]).append(auc)
            table.setdefault(algo,{}).setdefault(flag,{}).setdefault('prec',[]).append(prec)
            table.setdefault(algo,{}).setdefault(flag,{}).setdefault('rec',[]).append(rec)


for algo in table:
    if algo == 'jrip':
        continue
    print algo
    for flag in ['SOCIAL','PLACE','BOTH']:
        d = [algo]
        d.append(flag)
        auc = avg(table[algo][flag]['auc'])

        prec = avg(table[algo][flag]['prec'])
        rec = avg(table[algo][flag]['rec'])
        #print 'flag %s: %f - %f - %f'%(flag, auc,prec,rec)
        d.extend(map(lambda x: '%.2f'%x,(prec,rec,auc)))
        print ' & '.join(d), """\\\\"""

    print ''

for algo in bars:
    if algo == 'jrip':
        continue
    x = {}
    xerr = {}
    for month in bars[algo]:
        for flag in bars[algo][month]:
            x.setdefault(flag,[]).append(bars[algo][month][flag][0])
            xerr.setdefault(flag,[]).append(bars[algo][month][flag][-1])
            print algo, month, flag

    plt.figure()
    plt.clf()
    plt.axes(FIG_AXES2)
    x1 = x['SOCIAL']
    s1 = xerr['SOCIAL']
    x2 = x['PLACE']
    s2 = xerr['PLACE']
    x3 = x['BOTH']
    s3 = xerr['BOTH']
    print x1, x2, x3
    print s1,s2,s3
    ind = 1+np.arange(3)
    width = 0.2
    p1 = plt.bar(ind-width, x1, width, color='1.00',
            yerr=s1,ecolor='k')
    p2 = plt.bar(ind, x2, width, color='0.66',
            yerr=s2,ecolor='k')
    p3 = plt.bar(ind+width, x3, width, color='0.33',
            #yerr=s2,error_kw={'ecolor':'k'})
            yerr=s3,ecolor='k')

    plt.ylabel('AUC')
    plt.xticks(ind+width/2., ('Month 1', 'Month 2', 'Month 3'))
    plt.legend( (p1[0], p2[0], p3[0]), ('Social', 'Place', 'Place-social'),
            loc = 'lower left', prop={'size':8})
    plt.axis([0.5,3.5,0.75,1.0])
    plt.grid(True)
    plt.savefig('auc_bars_%s.pdf'%(algo))
    plt.close()
#
#for algo in bars:
#    if algo == 'jrip':
#        continue
#    x = {}
#    for month in bars[algo]:
#        for flag in bars[algo][month]:
#            x.setdefault(flag,[]).append(bars[algo][month][flag][1])
#            print algo, month, flag
#
#    plt.figure()
#    plt.clf()
#    plt.axes(FIG_AXES2)
#    x1 = x['SOCIAL']
#    x2 = x['PLACE']
#    x3 = x['BOTH']
#    print x1, x2, x3
#    ind = 1+np.arange(3)
#    width = 0.2
#    p1 = plt.bar(ind-width, x1, width, color='1.00')
#    p2 = plt.bar(ind, x2, width, color='0.50')
#    p3 = plt.bar(ind+width, x3, width, color='0.00')
#
#    plt.ylabel('Precision')
#    plt.xticks(ind+width/2., ('Month 1', 'Month 2', 'Month 3'))
#    plt.legend( (p1[0], p2[0], p3[0]), ('Social', 'Place', 'Place-social'),
#            loc = 'lower left', prop={'size':8})
#    plt.axis([0.5,3.5,0.0,1.0])
#    plt.grid(True)
#    plt.savefig('precision_bars_%s.pdf'%(algo))
#    plt.close()
#
#
#for algo in bars:
#    if algo == 'jrip':
#        continue
#    x = {}
#    for month in bars[algo]:
#        for flag in bars[algo][month]:
#            x.setdefault(flag,[]).append(bars[algo][month][flag][2])
#            print algo, month, flag
#
#    plt.figure()
#    plt.clf()
#    plt.axes(FIG_AXES2)
#    x1 = x['SOCIAL']
#    x2 = x['PLACE']
#    x3 = x['BOTH']
#    print x1, x2, x3
#    ind = 1+np.arange(3)
#    width = 0.2
#    p1 = plt.bar(ind-width, x1, width, color='1.00')
#    p2 = plt.bar(ind, x2, width, color='0.50')
#    p3 = plt.bar(ind+width, x3, width, color='0.00')
#
#    plt.ylabel('Recall')
#    plt.xticks(ind+width/2., ('Month 1', 'Month 2', 'Month 3'))
#    plt.legend( (p1[0], p2[0], p3[0]), ('Social', 'Place', 'Place-social'),
#            loc = 'lower left', prop={'size':8})
#    plt.axis([0.8,3.5,0.0,1.0])
#    plt.grid(True)
#    plt.savefig('recall_bars_%s.pdf'%(algo))
#    plt.close()
